from typing import List, Union
import torch
import torch.nn as nn
from torch import Tensor
from mmseg.models.segmentors import EncoderDecoder
from mmseg.registry import MODELS
from mmengine.structures import BaseDataElement


@MODELS.register_module()
class DualBranchSegmentor(EncoderDecoder):
    """EncoderDecoder with multiple decode heads (e.g. LCX + LAD)."""

    def _forward(self, inputs: Tensor, data_samples=None) -> List[Tensor]:
        """网络前向传播过程，返回双分支输出。
        
        Args:
            inputs (Tensor): 输入图像，形状为 (N, C, H, W)
            data_samples: 数据样本（在推理时未使用）
            
        Returns:
            List[Tensor]: 两个分支的输出列表 [lcx_output, lad_output]
        """
        x = self.extract_feat(inputs)
        
        # 双分支输出
        outputs = []
        for decode_head in self.decode_head:
            output = decode_head.forward(x)
            outputs.append(output)
        
        return outputs

    def encode_decode(self, inputs: Tensor, batch_img_metas: List[dict]) -> List[Tensor]:
        """编码解码过程，专门为双分支优化。
        
        Args:
            inputs (Tensor): 输入图像
            batch_img_metas (List[dict]): 图像元信息
            
        Returns:
            List[Tensor]: [lcx_logits, lad_logits] 双分支分割logits
        """
        x = self.extract_feat(inputs)
        
        # 双分支预测
        seg_logits = []
        for decode_head in self.decode_head:
            logits = decode_head.predict(x, batch_img_metas, self.test_cfg)
            seg_logits.append(logits)
        
        return seg_logits

    def postprocess_result(
        self,
        seg_logits: Union[torch.Tensor, List[torch.Tensor]],
        data_samples: List[BaseDataElement]
    ) -> List[BaseDataElement]:
        """后处理结果，支持双分支输出。"""
        if isinstance(seg_logits, (list, tuple)):
            # 多个 head
            for head_idx, logit in enumerate(seg_logits):
                batch_size, C, H, W = logit.shape
                preds = logit.argmax(dim=1)  # [B, H, W]
                if preds.max()!=0:
                    print('dual branch segmentor: ', int(preds.sum()))
                for i in range(batch_size):
                    # 不要写 dtype='tensor'，改成 torch.Tensor 或直接省略
                    data_samples[i].set_field(
                        preds[i].cpu(),
                        name=f'pred_seg_map_head{head_idx}',
                        dtype=torch.Tensor
                    )
        else:
            # 单个 head
            batch_size, C, H, W = seg_logits.shape
            preds = seg_logits.argmax(dim=1)
            for i in range(batch_size):
                data_samples[i].set_field(
                    preds[i].cpu(),
                    name='pred_seg_map',
                    dtype=torch.Tensor
                )
        return data_samples
